import re
import sys
import os

def get_version_from_pyproject(file_path):
    """从pyproject.toml文件中读取原始版本号"""
    try:
        with open(file_path, 'r') as file:
            content = file.read()
            match = re.search(r'version\s*=\s*["\']([0-9a-zA-Z.+-]+)["\']', content)
            if match:
                return match.group(1)
            else:
                print("未找到version信息", file=sys.stderr)
                return None
    except FileNotFoundError:
        print(f"错误: 文件 {file_path} 未找到", file=sys.stderr)
        return None
    except Exception as e:
        print(f"读取文件时发生错误: {str(e)}", file=sys.stderr)
        return None

def process_version(version):
    """处理版本号，保留前两位，最后一位用0代替"""
    parts = version.split('.')
    if len(parts) >= 3:
        processed_parts = parts[:2] + ['0']
        return '.'.join(processed_parts)
    else:
        print(f"警告: 版本号格式不符合预期 - {version}", file=sys.stderr)
        return version

if __name__ == "__main__":
    # 定义默认的pyproject.toml路径
    default_pyproject_path = "/workspace/omniinfer/pyproject.toml"
    
    # 检查命令行参数
    if len(sys.argv) == 2:
        # 只提供了输出路径，使用默认pyproject路径和默认文件名versions.txt
        output_dir = sys.argv[1]
        pyproject_path = default_pyproject_path
    elif len(sys.argv) == 3:
        # 提供了输出路径和pyproject路径
        output_dir = sys.argv[1]
        pyproject_path = sys.argv[2]
    else:
        print("读取参数错误")
        sys.exit(1)
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 构建完整输出文件路径（目录 + 默认文件名versions.txt）
    output_file_path = os.path.join(output_dir, "versions.txt")
    
    # 读取原始版本号
    original_version = get_version_from_pyproject(pyproject_path)
    
    if original_version:
        processed_version = process_version(original_version)
        
        # 将两个版本值写入指定文件
        with open(output_file_path, 'w') as f:
            f.write(f"ORIGINAL_VERSION={original_version}\n")
            f.write(f"PROCESSED_VERSION={processed_version}\n")
        
        print(f"版本信息已保存到: {output_file_path}")
        sys.exit(0)
    else:
        sys.exit(1)
    